def load_dataset(file_path):
    """加载数据集"""
    dataset = []
    with open(file_path, 'r') as file:
        for line in file.readlines():
            transaction = line.strip().split(',')
            dataset.append([frozenset([item]) for item in transaction])
    return dataset

def create_c1(dataset):
    """创建初始候选项集C1"""
    c1 = []
    for transaction in dataset:
        for item in transaction:
            if not [item] in c1:
                c1.append([item])
    c1.sort()
    return list(map(frozenset, c1))

def scan_d(d, candidates, min_support):
    """计算候选项集的支持度"""
    sscnt = {}
    for tid in d:
        for can in candidates:
            if can.issubset(tid):
                sscnt[can] = sscnt.get(can, 0) + 1
    num_items = float(len(d))
    retlist = []
    support_data = {}
    for key in sscnt:
        support = sscnt[key] / num_items
        if support >= min_support:
            retlist.insert(0, key)
        support_data[key] = support
    return retlist, support_data

def apriori_gen(lk, k):
    """连接步骤，生成Ck"""
    retlist = []
    lenlk = len(lk)
    for i in range(lenlk):
        for j in range(i+1, lenlk):
            l1 = list(lk[i])[:k-2]
            l2 = list(lk[j])[:k-2]
            l1.sort()
            l2.sort()
            if l1 == l2:
                retlist.append(lk[i] | lk[j])
    return retlist

def apriori(dataset, min_support=0.5):
    """Apriori算法主函数"""
    c1 = create_c1(dataset)
    d = list(map(set, dataset))
    l1, support_data = scan_d(d, c1, min_support)
    l = [l1]
    k = 2
    while (len(l[k-2]) > 0):
        ck = apriori_gen(l[k-2], k)
        lk, sup_k = scan_d(d, ck, min_support)
        support_data.update(sup_k)
        l.append(lk)
        k += 1
    return l, support_data

def generate_rules(freq_sets, support_data, min_confidence):
    """生成关联规则"""
    rules = []
    for i in range(1, len(freq_sets)):
        for freq_set in freq_sets[i]:
            m = len(freq_set)
            for ant in freq_sets[i-1]:
                if ant.issubset(freq_set):
                    conf = support_data[freq_set] / support_data[ant]
                    if conf >= min_confidence:
                        h = [list(freq_set - ant), list(ant), conf]
                        rules.append(h)
    return rules

# Load dataset
dataset = load_dataset('./wine/wine.data')

# Apply Apriori algorithm
l, support_data = apriori(dataset, min_support=0.02)

# Generate rules
rules = generate_rules(l, support_data, min_confidence=0.7)

# Print frequent itemsets
print("Frequent Itemsets:")
for level in l:
    for freq_set in level:
        print(freq_set)

# Print rules
print("\nAssociation Rules:")
for rule in rules:
    print(f"Rule: {rule[0]}, Antecedent: {rule[1]}, Confidence: {rule[2]}")